Puzzle 11 - flip解法

def flip_spec(a, out):
    for i in range(len(out)):
        out[i] = a[len(out) - i - 1]

def flip(a: TT["i"], i: int) -> TT["i"]:
    return a[(i-1-arange(i))]
    raise NotImplementedError


test_flip = make_test("flip", flip, flip_spec, add_sizes=["i"])

这道题要求将一维张量逆序翻转。根据目标规格：

$输出第k项为：out[k]=a[i-1-k]$

解法 ：倒序索引

构造反向索引向量：$i-1-arange(i)$，然后直接从 $a$ 中索引：

def flip(a: TT["i"], i: int) -> TT["i"]:
    return a[i - 1 - arange(i)]

运行示例（$i=4,a=[10,20,30,40]$）：

arange(4) 为 $[0,1,2,3]$

4 - 1 - arange(4) 得到 $[3,2,1,0]$

a[[3, 2, 1, 0]] 得到 $[40,30,20,10]$